"""
2D Navier-Stokes Solver

Implements the stream function-vorticity formulation for 2D incompressible flow
on L-shaped domains. Solves the coupled system:

1. Poisson equation: ∇²ψ = -ω (stream function from vorticity)
2. Vorticity transport: ∂ω/∂t + u·∇ω = (1/Re)∇²ω (pseudo-time to steady state)

where u = ∂ψ/∂y, v = -∂ψ/∂x are derived velocities.
"""

# Required numerical modules are already included in the main module

"""
    solve_lid_driven_2d_impl(options::LidDrivenOptions) -> LidDrivenResult

Core implementation of the 2D lid-driven cavity solver.

# Algorithm Overview
1. Initialize domain, grid, and fields
2. Create Poisson solver for stream function equation
3. Time-stepping loop:
   a. Solve Poisson equation: ∇²ψ = -ω
   b. Compute velocities: u = ∂ψ/∂y, v = -∂ψ/∂x  
   c. Update vorticity: explicit transport + boundary conditions
   d. Apply under-relaxation for stability
   e. Check convergence
4. Return solution and diagnostics

# Arguments
- `options::LidDrivenOptions`: Solver configuration

# Returns  
- `LidDrivenResult`: Complete solution with diagnostics
"""
function solve_lid_driven_2d_impl(options::LidDrivenOptions)
    
    if options.verbose
        println("🌀 Starting 2D Lid-Driven Cavity Solver")
        println("Configuration: N=$(options.n), Re=$(options.Re), solver=$(options.solver)")
    end
    
    # ================================================================
    # 1. Initialize domain and grid
    # ================================================================
    
    x, y, h = build_grid(options.n, Val(2))
    domain = create_lshape_domain(options.n, Val(2))
    nx, ny = length(x), length(y)
    
    if options.verbose
        print_domain_info(domain)
        print_grid_info(x, y, options.Re)
    end
    
    # ================================================================
    # 2. Initialize fields and solver
    # ================================================================
    
    # Primary fields
    ψ = zeros(Float64, nx, ny)      # Stream function
    ω = zeros(Float64, nx, ny)      # Vorticity
    u = zeros(Float64, nx, ny)      # x-velocity
    v = zeros(Float64, nx, ny)      # y-velocity
    
    # Set NaN in solid region for visualization
    ψ[.!domain.fluid_mask] .= NaN
    ω[.!domain.fluid_mask] .= NaN
    
    # Create Poisson solver
    poisson_solver = create_poisson_solver(domain, h, options)
    
    # Auto-adjust relaxation parameters
    α_ψ, α_ω = auto_adjust_relaxation_parameters(options.α_ψ, options.α_ω, 
                                                 options.n, options.solver, 
                                                 options.verbose)
    
    # ================================================================
    # 3. Time step and stability setup
    # ================================================================
    
    # Compute stable time step
    dt = compute_stable_timestep(h, options.Re, options.dt, options.verbose)
    
    # Initialize diagnostics
    convergence_history = Float64[]
    total_time = 0.0
    timing_breakdown = (poisson=0.0, velocity=0.0, vorticity=0.0, boundary=0.0)
    
    # ================================================================  
    # 4. Main time-stepping loop
    # ================================================================
    
    converged = false
    final_residual = Inf
    
    for step in 1:options.max_steps
        iter_start_time = time()
        
        # (a) Solve Poisson equation: ∇²ψ = -ω
        t0 = time()
        ψ_old = copy(ψ)
        
        poisson_result = solve_poisson!(poisson_solver, -ω, ψ)  # Note negative sign
        timing_breakdown = (timing_breakdown..., poisson = timing_breakdown.poisson + (time() - t0))
        
        if !poisson_result.converged && options.verbose
            @warn "Poisson solver did not converge at step $step"
        end
        
        # Apply relaxation to stream function
        under_relax!(ψ, ψ_old, α_ψ)
        
        # (b) Compute velocities from stream function
        t0 = time()
        compute_velocities_2d!(u, v, ψ, domain, h)
        apply_velocity_bc_2d!(u, v, domain, x, y, options.U_lid)
        timing_breakdown = (timing_breakdown..., velocity = timing_breakdown.velocity + (time() - t0))
        
        # (c) Adaptive time step based on CFL condition
        dt = update_adaptive_timestep(u, v, h, options.Re, dt, options.verbose && step % 50 == 0)
        
        # (d) Vorticity transport equation  
        t0 = time()
        ω_old = copy(ω)
        
        # Explicit vorticity transport: ∂ω/∂t + u·∇ω = (1/Re)∇²ω
        vorticity_transport_2d!(ω, u, v, domain, h, dt, options.Re)
        timing_breakdown = (timing_breakdown..., vorticity = timing_breakdown.vorticity + (time() - t0))
        
        # (e) Boundary conditions for vorticity
        t0 = time()
        apply_vorticity_bc_2d!(ω, ψ, domain, x, y, options.U_lid, h)
        timing_breakdown = (timing_breakdown..., boundary = timing_breakdown.boundary + (time() - t0))
        
        # Apply relaxation to vorticity
        under_relax!(ω, ω_old, α_ω)
        
        # (f) Convergence check
        Δω_norm = compute_convergence_norm(ω, ω_old, domain, options.n)
        push!(convergence_history, Δω_norm)
        
        # Check for convergence or divergence
        if Δω_norm < options.tol
            converged = true
            final_residual = Δω_norm
            if options.verbose
                println("✅ Converged at step $step: Δω = $(Printf.@sprintf("%.2e", Δω_norm))")
            end
            break
        elseif Δω_norm > 1e10 || isnan(Δω_norm)
            if options.verbose
                @warn "❌ Solution diverged at step $step: Δω = $Δω_norm"
            end
            final_residual = Δω_norm
            break
        end
        
        # Progress reporting
        iter_time = time() - iter_start_time
        total_time += iter_time
        
        if options.verbose && (step % 100 == 0 || step <= 10)
            u_max = maximum(abs.(u[domain.fluid_mask]))
            v_max = maximum(abs.(v[domain.fluid_mask]))  
            println(@sprintf("Step %4d: Δω=%.2e, dt=%.2e, |u|ₘₐₓ=%.3f, |v|ₘₐₓ=%.3f, t=%.3fs",
                           step, Δω_norm, dt, u_max, v_max, iter_time))
        end
    end
    
    # ================================================================
    # 5. Final processing and results
    # ================================================================
    
    final_step = min(length(convergence_history), options.max_steps)
    
    # Ensure boundary conditions are properly applied to final solution
    apply_boundary_conditions!(ψ, ω, u, v, domain, x, y, options.U_lid, h)
    
    if options.verbose
        print_timing_summary(timing_breakdown, total_time)
        print_solution_statistics(ψ, ω, u, v, domain, final_step)
    end
    
    # Create result structure
    result = LidDrivenResult(
        ψ = ψ,
        ω = ω, 
        u = u,
        v = v,
        w = nothing,  # 2D case
        x = x,
        y = y,
        z = nothing,  # 2D case
        converged = converged,
        iterations = final_step,
        residual_norm = final_residual,
        solve_time = total_time,
        convergence_history = options.save_history ? convergence_history : Float64[],
        options = options
    )
    
    return result
end

"""
    compute_velocities_2d!(u::Matrix{Float64}, v::Matrix{Float64}, 
                          ψ::Matrix{Float64}, domain::LShapeDomain{2}, h::Float64)

Compute velocity field from stream function using central differences.
u = ∂ψ/∂y, v = -∂ψ/∂x
"""
function compute_velocities_2d!(u::Matrix{Float64}, v::Matrix{Float64}, 
                               ψ::Matrix{Float64}, domain::LShapeDomain{2}, h::Float64)
    nx, ny = size(ψ)
    
    # Interior nodes: use central differences
    @inbounds for j in 2:ny-1, i in 2:nx-1
        if domain.fluid_mask[i, j]
            # u = ∂ψ/∂y (central difference in y-direction)
            ψ_n = domain.fluid_mask[i, j+1] ? ψ[i, j+1] : ψ[i, j]  # Mirror condition
            ψ_s = domain.fluid_mask[i, j-1] ? ψ[i, j-1] : ψ[i, j]
            u[i, j] = (ψ_n - ψ_s) / (2h)
            
            # v = -∂ψ/∂x (central difference in x-direction)  
            ψ_e = domain.fluid_mask[i+1, j] ? ψ[i+1, j] : ψ[i, j]
            ψ_w = domain.fluid_mask[i-1, j] ? ψ[i-1, j] : ψ[i, j]
            v[i, j] = -(ψ_e - ψ_w) / (2h)
        else
            u[i, j] = 0.0
            v[i, j] = 0.0
        end
    end
    
    # Handle boundary nodes (will be overwritten by boundary conditions)
    fill_boundary_velocities!(u, v, domain)
end

"""
    fill_boundary_velocities!(u::Matrix{Float64}, v::Matrix{Float64}, domain::LShapeDomain{2})

Fill boundary velocity values (will be corrected by boundary conditions later).
"""
function fill_boundary_velocities!(u::Matrix{Float64}, v::Matrix{Float64}, domain::LShapeDomain{2})
    nx, ny = size(u)
    
    # Set boundary and solid regions to zero (proper values applied later)
    @inbounds for j in 1:ny, i in 1:nx
        if domain.boundary_mask[i, j] || !domain.fluid_mask[i, j]
            u[i, j] = 0.0
            v[i, j] = 0.0
        end
    end
end

"""
    vorticity_transport_2d!(ω::Matrix{Float64}, u::Matrix{Float64}, v::Matrix{Float64},
                           domain::LShapeDomain{2}, h::Float64, dt::Float64, Re::Float64)

Update vorticity using explicit transport equation:
∂ω/∂t + u∇ω = (1/Re)∇²ω

Uses first-order upwind for convection and central differences for diffusion.
"""
function vorticity_transport_2d!(ω::Matrix{Float64}, u::Matrix{Float64}, v::Matrix{Float64},
                                 domain::LShapeDomain{2}, h::Float64, dt::Float64, Re::Float64)
    nx, ny = size(ω)
    h2 = h * h
    
    # Work with interior nodes only
    ω_new = copy(ω)
    
    @inbounds for j in 2:ny-1, i in 2:nx-1
        if !domain.interior_mask[i, j]
            continue  # Skip non-interior nodes
        end
        
        # Local velocity values
        u_c = u[i, j]  
        v_c = v[i, j]
        ω_c = ω[i, j]
        
        # Neighbor values with mirror conditions for solid boundaries
        ω_e = domain.fluid_mask[i+1, j] ? ω[i+1, j] : ω_c  
        ω_w = domain.fluid_mask[i-1, j] ? ω[i-1, j] : ω_c
        ω_n = domain.fluid_mask[i, j+1] ? ω[i, j+1] : ω_c
        ω_s = domain.fluid_mask[i, j-1] ? ω[i, j-1] : ω_c
        
        # Convection terms (first-order upwind)
        if u_c >= 0
            conv_x = u_c * (ω_c - ω_w) / h
        else
            conv_x = u_c * (ω_e - ω_c) / h
        end
        
        if v_c >= 0
            conv_y = v_c * (ω_c - ω_s) / h  
        else
            conv_y = v_c * (ω_n - ω_c) / h
        end
        
        convection = conv_x + conv_y
        
        # Diffusion term (central differences)
        diffusion = (ω_e + ω_w + ω_n + ω_s - 4*ω_c) / h2
        
        # Time update: ω^{n+1} = ω^n + dt*(-convection + diffusion/Re)
        ω_new[i, j] = ω_c + dt * (-convection + diffusion / Re)
    end
    
    # Update interior values
    @inbounds for j in 2:ny-1, i in 2:nx-1
        if domain.interior_mask[i, j]
            ω[i, j] = ω_new[i, j]
        end
    end
end

"""
    compute_stable_timestep(h::Float64, Re::Float64, dt_requested::Union{Float64,Nothing}, 
                           verbose::Bool) -> Float64

Compute stable time step based on CFL and diffusion constraints.
"""
function compute_stable_timestep(h::Float64, Re::Float64, dt_requested::Union{Float64,Nothing}, 
                                verbose::Bool)
    # Conservative stability factors
    cfl_factor = 0.45      # For explicit convection
    diffusion_factor = 0.14 # For explicit diffusion
    
    # Stability limits (assuming max velocity ~ 1)
    dt_cfl = cfl_factor * h
    dt_diffusion = diffusion_factor * h^2 * Re
    
    dt_stable = min(dt_cfl, dt_diffusion)
    
    if dt_requested === nothing || dt_requested > dt_stable
        if verbose
            old_str = dt_requested === nothing ? "auto" : @sprintf("%.2e", dt_requested)
            println("[Time step] Adjusted dt from $old_str to $(@sprintf("%.2e", dt_stable))")
            println("[Time step] Limited by: $(dt_diffusion < dt_cfl ? "diffusion" : "CFL")")
        end
        return dt_stable
    else
        return dt_requested
    end
end

"""
    update_adaptive_timestep(u::Matrix{Float64}, v::Matrix{Float64}, h::Float64, Re::Float64, 
                            dt_current::Float64, verbose::Bool) -> Float64

Update time step based on current velocity field.
"""  
function update_adaptive_timestep(u::Matrix{Float64}, v::Matrix{Float64}, h::Float64, Re::Float64, 
                                  dt_current::Float64, verbose::Bool)
    # Find maximum velocity magnitude
    u_max = maximum(abs.(u))
    v_max = maximum(abs.(v))
    vel_max = max(u_max, v_max, 1e-12)  # Avoid division by zero
    
    # CFL and diffusion constraints  
    cfl_factor = 0.45
    diffusion_factor = 0.14
    
    dt_cfl = cfl_factor * h / vel_max
    dt_diffusion = diffusion_factor * h^2 * Re
    
    dt_new = min(dt_cfl, dt_diffusion)
    
    if verbose && abs(dt_new - dt_current) / dt_current > 0.1
        println("[Adaptive dt] Changed from $(@sprintf("%.2e", dt_current)) to $(@sprintf("%.2e", dt_new)) (u_max=$(@sprintf("%.3e", vel_max)))")
    end
    
    return dt_new
end

"""
    compute_convergence_norm(ω_new::Matrix{Float64}, ω_old::Matrix{Float64}, 
                           domain::LShapeDomain{2}, n::Int) -> Float64

Compute convergence norm for vorticity field.
"""
function compute_convergence_norm(ω_new::Matrix{Float64}, ω_old::Matrix{Float64}, 
                                 domain::LShapeDomain{2}, n::Int)
    
    # Compute max change over interior nodes only
    max_change = 0.0
    
    @inbounds for idx in CartesianIndices(ω_new)
        if domain.interior_mask[idx]
            change = abs(ω_new[idx] - ω_old[idx])
            max_change = max(max_change, change)
        end
    end
    
    # Normalize by grid size for mesh-independent convergence
    return max_change / sqrt(n)
end

"""
    print_timing_summary(timing::NamedTuple, total_time::Float64)

Print breakdown of solver timing.
"""
function print_timing_summary(timing::NamedTuple, total_time::Float64)
    println("\n⏱️  Timing Summary:")
    println("  Poisson solves : $(@sprintf("%8.3f s (%4.1f%%)", timing.poisson, 100*timing.poisson/total_time))")
    println("  Velocity update: $(@sprintf("%8.3f s (%4.1f%%)", timing.velocity, 100*timing.velocity/total_time))")  
    println("  Vorticity trans: $(@sprintf("%8.3f s (%4.1f%%)", timing.vorticity, 100*timing.vorticity/total_time))")
    println("  Boundary conds : $(@sprintf("%8.3f s (%4.1f%%)", timing.boundary, 100*timing.boundary/total_time))")
    
    other_time = total_time - (timing.poisson + timing.velocity + timing.vorticity + timing.boundary)
    println("  Other          : $(@sprintf("%8.3f s (%4.1f%%)", other_time, 100*other_time/total_time))")
    println("  Total          : $(@sprintf("%8.3f s", total_time))")
end

"""
    print_solution_statistics(ψ::Matrix{Float64}, ω::Matrix{Float64}, 
                             u::Matrix{Float64}, v::Matrix{Float64}, 
                             domain::LShapeDomain{2}, iterations::Int)

Print statistics about the final solution.
"""
function print_solution_statistics(ψ::Matrix{Float64}, ω::Matrix{Float64}, 
                                  u::Matrix{Float64}, v::Matrix{Float64}, 
                                  domain::LShapeDomain{2}, iterations::Int)
    # Extract fluid region values
    ψ_fluid = ψ[domain.fluid_mask .& (.~isnan.(ψ))]
    ω_fluid = ω[domain.fluid_mask .& (.~isnan.(ω))]
    u_fluid = u[domain.fluid_mask]
    v_fluid = v[domain.fluid_mask]
    
    velocity_magnitude = sqrt.(u_fluid.^2 + v_fluid.^2)
    
    println("\n📊 Solution Statistics:")
    println("  Iterations: $iterations")
    println("  Stream function ψ: min=$(@sprintf("%.3e", minimum(ψ_fluid))), max=$(@sprintf("%.3e", maximum(ψ_fluid)))")
    println("  Vorticity ω      : min=$(@sprintf("%.3e", minimum(ω_fluid))), max=$(@sprintf("%.3e", maximum(ω_fluid)))")
    println("  Velocity |u|     : min=$(@sprintf("%.3e", minimum(velocity_magnitude))), max=$(@sprintf("%.3e", maximum(velocity_magnitude)))")
    println("  Mass conservation: ∫∇·u dV ≈ $(@sprintf("%.2e", compute_mass_conservation_error(u, v, domain)))")
end

"""
    compute_mass_conservation_error(u::Matrix{Float64}, v::Matrix{Float64}, domain::LShapeDomain{2}) -> Float64

Compute discrete mass conservation error ∫∇·u dV.
"""
function compute_mass_conservation_error(u::Matrix{Float64}, v::Matrix{Float64}, domain::LShapeDomain{2})
    nx, ny = size(u)
    h = 2.0 / (nx - 2)  # Approximate grid spacing
    
    div_u_sum = 0.0
    count = 0
    
    @inbounds for j in 2:ny-1, i in 2:nx-1
        if domain.interior_mask[i, j]
            # ∇·u = ∂u/∂x + ∂v/∂y (central differences)
            du_dx = (u[i+1, j] - u[i-1, j]) / (2h)
            dv_dy = (v[i, j+1] - v[i, j-1]) / (2h)
            div_u = du_dx + dv_dy
            
            div_u_sum += abs(div_u)
            count += 1
        end
    end
    
    return count > 0 ? div_u_sum / count : 0.0
end